#include "faceDetector.h"
#include <algorithm>
#include <iostream>

using namespace std;
FaceDetector::FaceDetector():
    _nms(0.4),
    _threshold(0.6),
    _mean_val{ 104.f, 117.f, 123.f },
    Net(new ncnn::Net())
{
}

inline void FaceDetector::Release() {
    if (Net != nullptr)
    {
        delete Net;
        Net = nullptr;
    }
}

FaceDetector::FaceDetector(const std::string &model_param, const std::string &model_bin) :
    _nms(0.4),
    _threshold(0.6),
    _mean_val{ 104.f, 117.f, 123.f },
    Net(new ncnn::Net())
{
    Init(model_param, model_bin);
}

void FaceDetector::Init(const std::string &model_param, const std::string &model_bin)
{
    int ret = Net->load_param(model_param.c_str());
    ret = Net->load_model(model_bin.c_str());
}
void FaceDetector::Detect(QImage& image, std::vector<Bbox>& boxes)
{
    const int max_side = 320;
    float long_side = fmax(image.width(), image.height());
    float scale = max_side / long_side;
    int new_w = int(image.width()*scale);
    int new_h = int(image.height()*scale);



    ncnn::Mat in = ncnn::Mat::from_pixels_resize(image.bits(), ncnn::Mat::PIXEL_RGBA2BGR, image.width(), image.height(), new_w, new_h);
    in.substract_mean_normalize(_mean_val, 0);
    ncnn::Extractor ex = Net->create_extractor();
    ex.set_light_mode(true);
    ex.set_num_threads(4);
    ex.input(0, in);
    ncnn::Mat out, out1, out2;

    // loc
    ex.extract("output0", out);

    // class
    ex.extract("530", out1);

    //landmark
    ex.extract("529", out2);

    std::vector<box> anchor;
    create_anchor(anchor, new_w, new_h);

    std::vector<Bbox> total_box;
    float *ptr = out.channel(0);
    float *ptr1 = out1.channel(0);
    float *landms = out2.channel(0);

    // #pragma omp parallel for num_threads(2)

    for (int i = 0; i < anchor.size(); ++i)
    {

        if (*(ptr1 + 1) > _threshold)
        {
            cout<<(*(ptr1 + 1))<<endl;

            box tmp = anchor[i];
            box tmp1;
            Bbox result;

            // loc and conf
            tmp1.cx = tmp.cx + *ptr * 0.1 * tmp.sx;
            tmp1.cy = tmp.cy + *(ptr + 1) * 0.1 * tmp.sy;
            tmp1.sx = tmp.sx * exp(*(ptr + 2) * 0.2);
            tmp1.sy = tmp.sy * exp(*(ptr + 3) * 0.2);

            result.x1 = (tmp1.cx - tmp1.sx / 2) * image.width();
            if (result.x1 < 0)
                result.x1 = 0;
            result.y1 = (tmp1.cy - tmp1.sy / 2) * image.height();
            if (result.y1 < 0)
                result.y1 = 0;
            result.x2 = (tmp1.cx + tmp1.sx / 2) * image.width();
            if (result.x2 > image.width())
                result.x2 = image.width();
            result.y2 = (tmp1.cy + tmp1.sy / 2)* image.height();
            if (result.y2 > image.height())
                result.y2 = image.height();
            result.s = *(ptr1 + 1);

            // landmark
            for (int j = 0; j < 5; ++j)
            {
                result.point[j]._x = (tmp.cx + *(landms + (j << 1)) * 0.1 * tmp.sx) * image.width();
                result.point[j]._y = (tmp.cy + *(landms + (j << 1) + 1) * 0.1 * tmp.sy) * image.height();
            }

            total_box.push_back(result);
        }
        ptr += 4;
        ptr1 += 2;
        landms += 10;
    }
    cout<<"total_box.size()"<<total_box.size()<<endl;

    std::sort(total_box.begin(), total_box.end(), cmp);
    nms(total_box, _nms);

    for (int j = 0; j < total_box.size(); ++j)
    {
        boxes.push_back(total_box[j]);
    }
}


void FaceDetector::Detect2(cv::Mat image, std::vector<Bbox>& boxes)
{
    const int max_side = 320;
    float long_side = fmax(image.cols, image.rows);
    float scale = max_side / long_side;
    int new_w = int(image.cols*scale);
    int new_h = int(image.rows*scale);



    ncnn::Mat in = ncnn::Mat::from_pixels_resize(image.data, ncnn::Mat::PIXEL_BGR, image.cols, image.rows, new_w, new_h);
    in.substract_mean_normalize(_mean_val, 0);
    ncnn::Extractor ex = Net->create_extractor();
    ex.set_light_mode(true);
    ex.set_num_threads(4);
    ex.input(0, in);
    ncnn::Mat out, out1, out2;

    // loc
    ex.extract("output0", out);

    // class
    ex.extract("530", out1);

    //landmark
    ex.extract("529", out2);

    std::vector<box> anchor;
    create_anchor(anchor, new_w, new_h);

    std::vector<Bbox> total_box;
    float *ptr = out.channel(0);
    float *ptr1 = out1.channel(0);
    float *landms = out2.channel(0);

    // #pragma omp parallel for num_threads(2)

    for (int i = 0; i < anchor.size(); ++i)
    {

        if (*(ptr1 + 1) > _threshold)
        {
            box tmp = anchor[i];
            box tmp1;
            Bbox result;

            // loc and conf
            tmp1.cx = tmp.cx + *ptr * 0.1 * tmp.sx;
            tmp1.cy = tmp.cy + *(ptr + 1) * 0.1 * tmp.sy;
            tmp1.sx = tmp.sx * exp(*(ptr + 2) * 0.2);
            tmp1.sy = tmp.sy * exp(*(ptr + 3) * 0.2);

            result.x1 = (tmp1.cx - tmp1.sx / 2) * image.cols;
            if (result.x1 < 0)
                result.x1 = 0;
            result.y1 = (tmp1.cy - tmp1.sy / 2) * image.rows;
            if (result.y1 < 0)
                result.y1 = 0;
            result.x2 = (tmp1.cx + tmp1.sx / 2) * image.cols;
            if (result.x2 > image.cols)
                result.x2 = image.cols;
            result.y2 = (tmp1.cy + tmp1.sy / 2)* image.rows;
            if (result.y2 > image.rows)
                result.y2 = image.rows;
            result.s = *(ptr1 + 1);

            // landmark
            for (int j = 0; j < 5; ++j)
            {
                result.point[j]._x = (tmp.cx + *(landms + (j << 1)) * 0.1 * tmp.sx) * image.rows;
                result.point[j]._y = (tmp.cy + *(landms + (j << 1) + 1) * 0.1 * tmp.sy) * image.rows;
            }

            total_box.push_back(result);
        }
        ptr += 4;
        ptr1 += 2;
        landms += 10;
    }

    std::sort(total_box.begin(), total_box.end(), cmp);
    nms(total_box, _nms);

    for (int j = 0; j < total_box.size(); ++j)
    {
        boxes.push_back(total_box[j]);
        if(boxes.size()==1)
            break;
    }
    cout<<"boxes_box.size()"<<boxes.size()<<endl;

}


inline bool FaceDetector::cmp(Bbox a, Bbox b) {
    if (a.s > b.s)
        return true;
    return false;
}

inline void FaceDetector::SetDefaultParams() {
    _nms = 0.4;
    _threshold = 0.6;
    _mean_val[0] = 104;
    _mean_val[1] = 117;
    _mean_val[2] = 123;
    Net = nullptr;

}

FaceDetector::~FaceDetector() {
    Release();
}

void FaceDetector::create_anchor(std::vector<box> &anchor, int w, int h)
{
    //    anchor.reserve(num_boxes);
    anchor.clear();
    std::vector<std::vector<int> > feature_map(4), min_sizes(4);
    float steps[] = { 8, 16, 32, 64 };
    for (int i = 0; i < feature_map.size(); ++i) {
        feature_map[i].push_back(ceil(h / steps[i]));
        feature_map[i].push_back(ceil(w / steps[i]));
    }
    std::vector<int> minsize1 = { 10, 16, 24 };
    min_sizes[0] = minsize1;
    std::vector<int> minsize2 = { 32, 48 };
    min_sizes[1] = minsize2;
    std::vector<int> minsize3 = { 64, 96 };
    min_sizes[2] = minsize3;
    std::vector<int> minsize4 = { 128, 192, 256 };
    min_sizes[3] = minsize4;


    for (int k = 0; k < feature_map.size(); ++k)
    {
        std::vector<int> min_size = min_sizes[k];
        for (int i = 0; i < feature_map[k][0]; ++i)
        {
            for (int j = 0; j < feature_map[k][1]; ++j)
            {
                for (int l = 0; l < min_size.size(); ++l)
                {
                    float s_kx = min_size[l] * 1.0 / w;
                    float s_ky = min_size[l] * 1.0 / h;
                    float cx = (j + 0.5) * steps[k] / w;
                    float cy = (i + 0.5) * steps[k] / h;
                    box axil = { cx, cy, s_kx, s_ky };
                    anchor.push_back(axil);
                }
            }
        }

    }

}

void FaceDetector::nms(std::vector<Bbox> &input_boxes, float NMS_THRESH)
{
    std::vector<float>vArea(input_boxes.size());
    for (int i = 0; i < int(input_boxes.size()); ++i)
    {
        vArea[i] = (input_boxes.at(i).x2 - input_boxes.at(i).x1 + 1)
            * (input_boxes.at(i).y2 - input_boxes.at(i).y1 + 1);
    }
    for (int i = 0; i < int(input_boxes.size()); ++i)
    {
        for (int j = i + 1; j < int(input_boxes.size());)
        {
            float xx1 = fmax(input_boxes[i].x1, input_boxes[j].x1);
            float yy1 = fmax(input_boxes[i].y1, input_boxes[j].y1);
            float xx2 = fmin(input_boxes[i].x2, input_boxes[j].x2);
            float yy2 = fmin(input_boxes[i].y2, input_boxes[j].y2);
            float w = fmax(float(0), xx2 - xx1 + 1);
            float   h = fmax(float(0), yy2 - yy1 + 1);
            float   inter = w * h;
            float ovr = inter / (vArea[i] + vArea[j] - inter);
            if (ovr >= NMS_THRESH)
            {
                input_boxes.erase(input_boxes.begin() + j);
                vArea.erase(vArea.begin() + j);
            }
            else
            {
                j++;
            }
        }
    }
}
